package io.github.openfeign.spring.enhancer;

import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.NotFoundException;
import javassist.bytecode.BadBytecode;
import javassist.bytecode.CodeAttribute;
import javassist.bytecode.CodeIterator;
import javassist.bytecode.ConstPool;
import javassist.bytecode.LineNumberAttribute;
import javassist.bytecode.MethodInfo;
import javassist.bytecode.Opcode;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.function.BiFunction;

/**
 * 抽象字节码转换器，提供基于Javassist的字节码操作基础能力
 *
 * <p>实现以下核心功能：
 * <ul>
 *   <li>定位字节码中的特定指令位置</li>
 *   <li>在指定位置插入自定义代码</li>
 *   <li>持久化修改后的字节码</li>
 * </ul>
 *
 * <p>子类需实现具体转换逻辑，应用于Feign客户端的multipart流式处理增强
 */
@Slf4j
public abstract class AbstractBytecodeTransformer {

    /**
     * Javassist类池，用于操作类字节码
     */
    final ClassPool classPool;

    /**
     * 修改后的类文件输出路径（用于调试）
     */
    final String classOutPath;

    /**
     * 初始化字节码转换器
     *
     * @param classPool    类池实例，用于字节码操作
     * @param classOutPath 修改后的类文件输出路径（可空）
     */
    public AbstractBytecodeTransformer(final ClassPool classPool, final String classOutPath) {
        this.classPool = classPool;
        this.classOutPath = classOutPath;
    }

    /**
     * 执行字节码转换流程
     *
     * <p>封装异常处理，将检查异常转换为运行时异常
     */
    public void process() {
        try {
            transformBytecode();
        } catch (final NotFoundException | CannotCompileException | IOException | BadBytecode e) {
            throw new RuntimeException("Bytecode transformation failed", e);
        }
    }

    /**
     * 抽象转换方法，由子类实现具体字节码修改逻辑
     *
     * @throws NotFoundException      当目标类/方法不存在时抛出
     * @throws CannotCompileException 代码插入导致编译错误时抛出
     * @throws java.io.IOException            类文件读写异常时抛出
     * @throws BadBytecode            遇到非法字节码结构时抛出
     */
    abstract void transformBytecode() throws NotFoundException, CannotCompileException, BadBytecode, IOException;

    /* 字节码指令定位方法组 */

    /**
     * 定位IFNULL指令位置
     *
     * @param iter  字节码指令迭代器
     * @param start 起始搜索位置
     * @return IFNULL指令的字节码位置（未找到返回-1）
     */
    static int locateIfNullInstruction(final CodeIterator iter, final int start) {
        return locateNextOpcodePosition(iter, start, Opcode.IFNULL);
    }

    /**
     * 通用指令定位方法
     *
     * @param iter     字节码指令迭代器
     * @param startPos 起始搜索位置
     * @param targetOp 目标操作码
     * @return 目标操作码的位置（未找到返回-1）
     * @throws IllegalStateException 当字节码解析异常时抛出
     */
    static int locateNextOpcodePosition(final CodeIterator iter, final int startPos, final int targetOp) {
        iter.move(startPos);
        try {
            while (iter.hasNext()) {
                final int currentPos = iter.next();
                if (iter.byteAt(currentPos) == targetOp) {
                    return currentPos;
                }
            }
            return -1;
        } catch (final BadBytecode e) {
            throw new IllegalStateException("Bytecode parsing error", e);
        }
    }

    /* 代码插入方法组 */

    /**
     * 在指定方法调用位置插入代码片段
     *
     * @param targetMethod    目标方法
     * @param className       目标类全限定名
     * @param methodName      目标方法名称
     * @param codeSnippet     待插入的代码片段
     * @param triggerOpcode   触发插入的操作码（如方法调用指令）
     * @param positionLocator 位置定位策略函数
     * @throws BadBytecode            字节码解析异常
     * @throws CannotCompileException 代码插入导致编译错误
     */
    void insertCodeAtTargetMethod(final CtMethod targetMethod,
                                  final String className,
                                  final String methodName,
                                  final String codeSnippet,
                                  final int triggerOpcode,
                                  final BiFunction<CodeIterator, Integer, Integer> positionLocator)
            throws BadBytecode, CannotCompileException {

        final MethodInfo methodInfo = targetMethod.getMethodInfo();
        final CodeAttribute codeAttr = methodInfo.getCodeAttribute();
        final CodeIterator iterator = codeAttr.iterator();
        final LineNumberAttribute lineNumAttr = (LineNumberAttribute) codeAttr.getAttribute(LineNumberAttribute.tag);
        final ConstPool constPool = methodInfo.getConstPool();

        while (iterator.hasNext()) {
            final int pos = iterator.next();
            if (iterator.byteAt(pos) == triggerOpcode) {
                final int methodRefIndex = iterator.u16bitAt(pos + 1);
                if (isTargetMethod(constPool, methodRefIndex, className, methodName)) {
                    final int insertionLine = resolveLineNumber(lineNumAttr, positionLocator.apply(iterator, pos));
                    if (insertionLine != -1) {
                        injectCodeAtLine(targetMethod, insertionLine, codeSnippet);
                    }
                }
            }
        }
    }

    /**
     * 在方法的最后return语句前插入代码
     *
     * @param targetMethod 目标方法
     * @param returnOpcode return操作码（如ARETURN）
     * @param codeSnippet  待插入的代码
     * @throws BadBytecode            字节码解析异常
     * @throws CannotCompileException 代码插入导致编译错误
     */
    void insertBeforeFinalReturn(final CtMethod targetMethod,
                                 final int returnOpcode,
                                 final String codeSnippet)
            throws BadBytecode, CannotCompileException {

        final MethodInfo methodInfo = targetMethod.getMethodInfo();
        final CodeAttribute codeAttr = methodInfo.getCodeAttribute();
        final CodeIterator iterator = codeAttr.iterator();
        final LineNumberAttribute lineNumAttr = (LineNumberAttribute) codeAttr.getAttribute(LineNumberAttribute.tag);

        final List<Integer> returnPoints = new ArrayList<>();
        while (iterator.hasNext()) {
            final int pos = iterator.next();
            if (iterator.byteAt(pos) == returnOpcode) {
                returnPoints.add(resolveLineNumber(lineNumAttr, pos));
            }
        }

        if (returnPoints.isEmpty()) {
            return;
        }
        final int lastReturnLine = returnPoints.stream().max(Integer::compareTo).get();
        injectCodeAtLine(targetMethod, lastReturnLine, codeSnippet);
    }

    /**
     * 在目标方法的指定行号注入代码片段
     *
     * @param targetMethod 需要增强的CtMethod对象
     * @param lineNumber   目标源代码行号（基于调试符号表）
     * @param codeToInject 要注入的Java代码字符串
     * @throws CannotCompileException 当出现以下情况时抛出：
     *                                <ul>
     *                                  <li>插入位置超出方法行号范围</li>
     *                                  <li>注入代码存在语法错误</li>
     *                                  <li>字节码验证失败</li>
     *                                </ul>
     */
    private void injectCodeAtLine(final CtMethod targetMethod,
                                  final int lineNumber,
                                  final String codeToInject) throws CannotCompileException {

        if (log.isDebugEnabled()) {
            log.debug(
                "Injecting code at line {} [Transformer: {}]\nCode content:\n{}",
                lineNumber,
                this.getClass().getSimpleName(),
                codeToInject
            );
        }

        targetMethod.insertAt(lineNumber, codeToInject);
    }


    // 辅助方法

    /**
     * 转换字节码位置到源代码行号
     *
     * @param lineNumAttr 行号表属性
     * @param bytecodePos 字节码指令位置
     * @return 对应的源代码行号（无法解析返回-1）
     */
    static int resolveLineNumber(final LineNumberAttribute lineNumAttr, final int bytecodePos) {
        return (bytecodePos != -1 && lineNumAttr != null) ? lineNumAttr.toLineNumber(bytecodePos) : -1;
    }

    /**
     * 持久化修改后的类文件
     *
     * @param modifiedClass 已修改的CtClass实例
     */
    void persistModifiedClass(final CtClass modifiedClass) {
        if (StringUtils.isBlank(classOutPath)){
            return;
        }
        try {
            modifiedClass.writeFile(classOutPath);
        } catch (CannotCompileException | IOException e) {
            log.warn("Failed to persist modified class: {}", modifiedClass.getName(), e);
        }
    }

    /**
     * 将修改后的类重新定义到JVM
     *
     * @param modifiedClass 已修改的CtClass实例
     * @throws CannotCompileException 类重定义失败时抛出
     */
    void redefineClassInMemory(final CtClass modifiedClass) throws CannotCompileException {
        persistModifiedClass(modifiedClass);
        modifiedClass.toClass();
    }

    /**
     * 验证常量池中的方法引用
     *
     * @param constPool      常量池实例
     * @param refIndex       方法引用索引
     * @param expectedClass  预期类名
     * @param expectedMethod 预期方法名
     */
    private boolean isTargetMethod(final ConstPool constPool,
                                   final int refIndex,
                                   final String expectedClass,
                                   final String expectedMethod) {
        return expectedClass.equals(constPool.getMethodrefClassName(refIndex))
                && expectedMethod.equals(constPool.getMethodrefName(refIndex));
    }
}
